﻿using System;
using System.Buffers;
using System.Collections.Generic;
using System.Diagnostics;
using System.Diagnostics.CodeAnalysis;
using System.IO;
using System.Linq;
using System.Runtime.CompilerServices;
using System.Runtime.InteropServices;
using System.Threading;
using System.Threading.Tasks;
using SixLabors.ImageSharp;
using SixLabors.ImageSharp.Formats;
using SixLabors.ImageSharp.Formats.Bmp;
using SixLabors.ImageSharp.Formats.Jpeg;
using SixLabors.ImageSharp.Formats.Png;
using SixLabors.ImageSharp.PixelFormats;
using SixLabors.ImageSharp.Processing;
using WholeSlideImageLibrary.DeepZoom;
using WholeSlideImageLibrary.Features;

namespace WholeSlideImageLibrary.Formats.Dmetrix
{
    internal sealed class DmetrixWholeSlideImage : WholeSlideImage, IDeepZoomProfileHintFeature, IAssociatedImagesFeature, IAssociatedImagesEncodingFeature
    {
        private FileSourceReaderCache? _source;
        private readonly DmetrixFileHeader _header;
        internal int _firstValidLayer;

        internal DmetrixWholeSlideImage(FileSourceInfo fileSource, in DmetrixFileHeader header)
        {
            _source = new FileSourceReaderCache(fileSource);
            _header = header;
            _firstValidLayer = FindFirstValidLayer(in _header);
        }

        private static int FindFirstValidLayer(in DmetrixFileHeader header)
        {
            DmetrixFileLayerIndex[]? layerIndexes = header.LayerIndexes;
            for (int i = 0; i < header.MaxLayer; i++)
            {
                if (layerIndexes[i].LayerStartPosition != 0)
                {
                    return i;
                }
            }
            return -1;
        }


        #region Properties

        /// <summary>
        /// Whether this file is encrypted.
        /// </summary>
        public bool IsEncrypted => _header.IsEncrypted;

        /// <summary>
        /// Sanner device. (10 bytes)
        /// </summary>
        public string Device => _header.Device;

        /// <summary>
        /// Craetion date.
        /// </summary>
        public DateTime Date => _header.Date;

        /// <summary>
        /// Image width.
        /// </summary>
        public override int Width => _header.Width;

        /// <summary>
        /// Image height.
        /// </summary>
        public override int Height => _header.Height;

        /// <summary>
        /// Total layer count.
        /// </summary>
        internal int LayerCount => _header.MaxLayer;

        /// <summary>
        /// Microns per pixel in X direction.
        /// </summary>
        public override double MicronsPerPixelX => _header.XLength;

        /// <summary>
        /// Microns per pixel in Y direction.
        /// </summary>
        public override double MicronsPerPixelY => _header.YLength;

        /// <summary>
        /// Maximum magnification  of this image.
        /// </summary>
        public int Magnification => _header.Multple;

        internal int FirstValidLayer => _firstValidLayer;

        public override int LevelCount => _header.MaxLayer - _firstValidLayer;

        #endregion

        internal DeepZoomLayer[]? _layers;

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        internal DeepZoomLayer[] EnsureDeepZoomLayersInitialized()
        {
            return _layers is null ? _layers = DeepZoomHelper.CalculateDeepZoomLayers(Width, Height) : _layers;
        }


        #region DeepZoom dimensions

        public override System.Drawing.Size GetLevelSize(int level)
        {
            EnsureNotDisposed();

            int levelCount = LevelCount;
            if ((uint)level > (uint)levelCount)
            {
                throw new ArgumentOutOfRangeException(nameof(level));
            }

            level = levelCount - 1 - level;
            int layer = level + _firstValidLayer;

            DeepZoomLayer[]? layers = EnsureDeepZoomLayersInitialized();
            return new System.Drawing.Size(layers[layer].Width, layers[layer].Height);
        }

        #endregion

        #region Get Tile

        public async Task<DmetrixFileTileIndex?> FindTileIndexAsync(FileSourceReader reader, int layer, int row, int column, CancellationToken cancellationToken)
        {
            byte[] buffer = ArrayPool<byte>.Shared.Rent(DmetrixFileTileIndex.MarshalStructureSize);
            try
            {
                return await FindTileIndexWorkspaceAsync(reader, layer, row, column, buffer, cancellationToken).ConfigureAwait(false);
            }
            finally
            {
                ArrayPool<byte>.Shared.Return(buffer);
            }
        }

        private async Task<DmetrixFileTileIndex?> FindTileIndexWorkspaceAsync(FileSourceReader reader, int layer, int row, int column, Memory<byte> workspace, CancellationToken cancellationToken)
        {
            Debug.Assert(workspace.Length >= DmetrixFileTileIndex.MarshalStructureSize);

            long tileIndexOffset = FindTileIndexOffset(layer, row, column);
            DmetrixFileTileIndex tileIndex;

            byte[] buffer = ArrayPool<byte>.Shared.Rent(DmetrixFileTileIndex.MarshalStructureSize);
            try
            {
                int bytesRead = await reader.ReadAsync(tileIndexOffset, buffer.AsMemory(0, DmetrixFileTileIndex.MarshalStructureSize), cancellationToken).ConfigureAwait(false);

                if (!DmetrixFileTileIndex.TryParse(buffer.AsSpan(0, bytesRead), out tileIndex))
                {
                    return default;
                }
            }
            finally
            {
                ArrayPool<byte>.Shared.Return(buffer);
            }

            return (layer == tileIndex.Layer && row == tileIndex.Row && column == tileIndex.Column) ? tileIndex : default(DmetrixFileTileIndex?);
        }

        private long FindTileIndexOffset(int layer, int row, int column)
        {
            long tileIndexOffset;
            switch (layer)
            {
                case -2:
                    tileIndexOffset = 370;
                    break;
                case -1:
                    tileIndexOffset = 348;
                    break;
                default:
                    DmetrixFileLayerIndex index = _header.LayerIndexes[layer];
                    tileIndexOffset = index.LayerStartPosition + ((long)column * (index.MaxRow + 1) + row) * DmetrixFileTileIndex.MarshalStructureSize;
                    break;
            }
            return tileIndexOffset;
        }

        #endregion

        #region Read Region

        public override async ValueTask ReadRegionAsync<T>(int level, int x, int y, int width, int height, Memory<T> buffer, CancellationToken cancellationToken)
        {
            FileSourceReaderCache? source = EnsureNotDisposed();

            int levelCount = LevelCount;
            if ((uint)level > (uint)levelCount)
            {
                throw new ArgumentOutOfRangeException(nameof(level));
            }
            if (Unsafe.SizeOf<T>() * buffer.Length < 4 * width * height)
            {
                throw new ArgumentException("destination is too small.", nameof(buffer));
            }

            level = levelCount - 1 - level;
            int layer = level + _firstValidLayer;

            const int tileSize = 256; // Dmetrix文件必是256

            DeepZoomLayer[]? layers = EnsureDeepZoomLayersInitialized();

            int colStart = (int)(x / tileSize);
            int rowStart = (int)(y / tileSize);
            int colEnd = (int)((x + width - 1) / tileSize);
            int rowEnd = (int)((y + height - 1) / tileSize);

            int layerColCount = (layers[layer].Width + (tileSize - 1)) / tileSize;
            int layerRowCount = (layers[layer].Height + (tileSize - 1)) / tileSize;
            if (colEnd >= layerColCount)
            {
                colEnd = layerColCount - 1;
            }
            if (rowEnd >= layerRowCount)
            {
                rowEnd = layerRowCount - 1;
            }

            int colLength = colEnd - colStart + 1;
            int rowLength = rowEnd - rowStart + 1;

            using var destination = Image.WrapMemory(buffer.Cast<T, Rgba32>(), width, height);

            const int minimumBufferSize = DmetrixFileTileIndex.MarshalStructureSize;
            int headerSize = _header.HSize;
            byte[]? pooledBuffer = null;
            FileSourceReader? reader = null;
            try
            {
                reader = await source.RentAsync(cancellationToken).ConfigureAwait(false);

                using (var canvas = new Image<Rgba32>(colLength * tileSize, rowLength * tileSize))
                {
                    // 读取所有覆盖的patch
                    for (int col = 0; col < colLength; col++)
                    {
                        for (int row = 0; row < rowLength; row++)
                        {
                            ExpandRentedArray(ref pooledBuffer, minimumBufferSize);
                            DmetrixFileTileIndex? tileIndex = await FindTileIndexWorkspaceAsync(reader, layer, rowStart + row, colStart + col, pooledBuffer, cancellationToken).ConfigureAwait(false);
                            if (tileIndex == null)
                            {
                                throw new InvalidDataException("Tile not found.");
                            }

                            int size = tileIndex.GetValueOrDefault().Size;
                            ExpandRentedArray(ref pooledBuffer, Math.Max(minimumBufferSize, size));
                            int bytesRead = await reader.ReadAsync(headerSize + tileIndex.GetValueOrDefault().StartPosition, pooledBuffer.AsMemory(0, size), cancellationToken).ConfigureAwait(false);
                            if (bytesRead != size)
                            {
                                throw new InvalidDataException();
                            }

                            using var tile = Image.Load<Rgba32>(new MemoryStream(pooledBuffer, 0, size));
                            canvas.Mutate(ctx =>
                            {
                                ctx.DrawImage(tile, new SixLabors.ImageSharp.Point(col * tileSize, row * tileSize), opacity: 1);
                            });
                        }
                    }

                    // 剪裁
                    destination.Mutate(ctx =>
                    {
                        ctx.DrawImage(canvas, new SixLabors.ImageSharp.Point(-(int)(x - colStart * tileSize), -(int)(y - rowStart * tileSize)), opacity: 1);
                    });
                }
            }
            finally
            {
                if (!(pooledBuffer is null))
                {
                    ArrayPool<byte>.Shared.Return(pooledBuffer);
                }
                if (!(reader is null))
                {
                    await source.ReturnAsync(reader).ConfigureAwait(false);
                }
            }
        }

        #endregion

        #region Associated Images

        private DmetrixAssociatedImageInfo[]? _associatedImage;
        private TaskCompletionSource<object?> _associatedImagesLoaded = new TaskCompletionSource<object?>(TaskCreationOptions.RunContinuationsAsynchronously);
        private int _associatedImagesLock = 1;

        private ValueTask EnsureAssociatedImagesLoaded()
        {
            if (!(_associatedImage is null))
            {
                return default;
            }

            return new ValueTask(StartLoadAssociatedImages());
        }

        private Task StartLoadAssociatedImages()
        {
            FileSourceReaderCache? cache = EnsureNotDisposed();

            bool lockTaken = Interlocked.Exchange(ref _associatedImagesLock, 0) == 1;
            if (lockTaken)
            {
                _ = Task.Run(async () =>
                {
                    try
                    {
                        await LoadAssociatedImageAsync(cache).ConfigureAwait(false);
                    }
                    catch (Exception)
                    {
                        // Do nothing
                    }
                });
            }

            return _associatedImagesLoaded.Task;
        }

        private async Task LoadAssociatedImageAsync(FileSourceReaderCache cache)
        {
            long offset;
            int size, bytesRead;
            int headerSize = _header.HSize;
            byte[] buffer = ArrayPool<byte>.Shared.Rent(DmetrixFileTileIndex.MarshalStructureSize);
            FileSourceReader? reader = null;
            try
            {
                reader = await cache.RentAsync(default).ConfigureAwait(false);

                IImageInfo? imageInfo;
                IImageFormat? format;

                // label
                DmetrixFileTileIndex? tileIndex = await FindTileIndexWorkspaceAsync(reader, -1, 0, 0, buffer, CancellationToken.None).ConfigureAwait(false);
                DmetrixAssociatedImageInfo? label = null;
                try
                {
                    if (tileIndex.HasValue)
                    {
                        offset = headerSize + tileIndex.GetValueOrDefault().StartPosition;
                        size = tileIndex.GetValueOrDefault().Size;
                        ExpandRentedArray(ref buffer, size);
                        bytesRead = await reader.ReadAsync(offset, buffer.AsMemory(0, size), CancellationToken.None).ConfigureAwait(false);
                        if (bytesRead != size)
                        {
                            throw new InvalidDataException();
                        }

                        (imageInfo, format) = await Image.IdentifyWithFormatAsync(new MemoryStream(buffer, 0, size)).ConfigureAwait(false);
                        EncodedImageFileFormat fileFormat = format switch
                        {
                            JpegFormat _ => EncodedImageFileFormat.Jpeg,
                            BmpFormat _ => EncodedImageFileFormat.Bmp,
                            PngFormat _ => EncodedImageFileFormat.Png,
                            _ => EncodedImageFileFormat.None,
                        };

                        label = new DmetrixAssociatedImageInfo("label", imageInfo.Width, imageInfo.Height, offset, size, fileFormat);
                    }
                }
                catch (Exception)
                {
                    // Do nothing
                }

                // thumbnail
                tileIndex = await FindTileIndexWorkspaceAsync(reader, -2, 0, 0, buffer, CancellationToken.None).ConfigureAwait(false);
                DmetrixAssociatedImageInfo? thumbnail = null;
                try
                {
                    if (tileIndex.HasValue)
                    {
                        offset = headerSize + tileIndex.GetValueOrDefault().StartPosition;
                        size = tileIndex.GetValueOrDefault().Size;
                        ExpandRentedArray(ref buffer, size);
                        bytesRead = await reader.ReadAsync(offset, buffer.AsMemory(0, size), CancellationToken.None).ConfigureAwait(false);
                        if (bytesRead != size)
                        {
                            throw new InvalidDataException();
                        }

                        (imageInfo, format) = await Image.IdentifyWithFormatAsync(new MemoryStream(buffer, 0, size)).ConfigureAwait(false);
                        EncodedImageFileFormat fileFormat = format switch
                        {
                            JpegFormat _ => EncodedImageFileFormat.Jpeg,
                            BmpFormat _ => EncodedImageFileFormat.Bmp,
                            PngFormat _ => EncodedImageFileFormat.Png,
                            _ => EncodedImageFileFormat.None,
                        };

                        thumbnail = new DmetrixAssociatedImageInfo("thumbnail", imageInfo.Width, imageInfo.Height, offset, size, fileFormat);
                    }
                }
                catch (Exception)
                {
                    // Do nothing
                }

                _associatedImage = (new DmetrixAssociatedImageInfo?[] { label, thumbnail }).Where(i => !(i is null)).ToArray()!;
            }
            finally
            {
                _associatedImagesLoaded.TrySetResult(null);

                ArrayPool<byte>.Shared.Return(buffer);
                if (!(reader is null))
                {
                    await reader.DisposeAsync().ConfigureAwait(false);
                }
            }
        }

        private DmetrixAssociatedImageInfo? FindAssociatedImageInfo(string name)
        {
            Debug.Assert(_associatedImage != null);
            if (name is null)
            {
                return null;
            }
            foreach (DmetrixAssociatedImageInfo? item in _associatedImage!)
            {
                if (item.Name.Equals(name, StringComparison.Ordinal))
                {
                    return item;
                }
            }
            return null;
        }

        public async ValueTask<IEnumerable<string>> GetAllAssociatedImagesAsync(CancellationToken cancellationToken = default)
        {
            await EnsureAssociatedImagesLoaded().ConfigureAwait(false);

            return _associatedImage.Select(i => i.Name)!;
        }

        public async ValueTask<System.Drawing.Size> GetAssociatedImageSizeAsync(string name, CancellationToken cancellationToken = default)
        {
            if (name is null)
            {
                throw new ArgumentNullException(nameof(name));
            }

            await EnsureAssociatedImagesLoaded().ConfigureAwait(false);

            DmetrixAssociatedImageInfo? image = FindAssociatedImageInfo(name);
            if (image is null)
            {
                throw new KeyNotFoundException();
            }

            return new System.Drawing.Size(image.Width, image.Height);
        }

        public async ValueTask ReadAssociatedImageAsync<T>(string name, Memory<T> buffer, CancellationToken cancellationToken = default) where T : unmanaged
        {
            FileSourceReaderCache? cache = EnsureNotDisposed();
            await EnsureAssociatedImagesLoaded().ConfigureAwait(false);

            if (name is null)
            {
                throw new ArgumentNullException(nameof(name));
            }

            DmetrixAssociatedImageInfo? image = FindAssociatedImageInfo(name);
            if (image is null)
            {
                throw new KeyNotFoundException();
            }

            if (Unsafe.SizeOf<T>() * buffer.Length < 4 * image.Width * image.Height)
            {
                throw new ArgumentException("destination is too small.", nameof(buffer));
            }

            byte[] pooledBuffer = ArrayPool<byte>.Shared.Rent(image.Length);
            FileSourceReader? reader = null;
            try
            {
                reader = await cache.RentAsync(cancellationToken).ConfigureAwait(false);

                int bytesRead = await reader.ReadAsync(image.Offset, pooledBuffer.AsMemory(0, image.Length), cancellationToken).ConfigureAwait(false);
                if (bytesRead != image.Length)
                {
                    throw new InvalidDataException();
                }

                using var canvas = Image.Load<Rgba32>(new MemoryStream(pooledBuffer, 0, image.Length));
                if (canvas.Width != image.Width || canvas.Height != image.Height)
                {
                    throw new InvalidDataException();
                }

                CopyPixels(canvas, buffer);
            }
            finally
            {
                ArrayPool<byte>.Shared.Return(pooledBuffer);
                if (!(reader is null))
                {
                    await reader.DisposeAsync().ConfigureAwait(false);
                }
            }
        }

        private static void CopyPixels<T>(Image<Rgba32> source, Memory<T> destination) where T : unmanaged
        {
            Debug.Assert(Unsafe.SizeOf<T>() * destination.Length >= 4 * source.Width * source.Height);

            Span<byte> destinationSpan = MemoryMarshal.AsBytes(destination.Span);

            if (source.TryGetSinglePixelSpan(out Span<Rgba32> sourceSpan))
            {
                MemoryMarshal.AsBytes(sourceSpan).CopyTo(destinationSpan.Slice(0, 4 * sourceSpan.Length));
            }
            else
            {
                int height = source.Height;
                int rowByteCount = 4 * source.Width;
                for (int i = 0; i < height; i++)
                {
                    Span<Rgba32> rowSpan = source.GetPixelRowSpan(i);
                    MemoryMarshal.AsBytes(rowSpan).CopyTo(destinationSpan.Slice(0, rowByteCount));
                    destinationSpan = destinationSpan.Slice(rowByteCount);
                }
            }
        }

        public async ValueTask<EncodedImageFileFormat> GetAssociatedImageEncodingSupportAsync(string name, CancellationToken cancellationToken = default)
        {
            await EnsureAssociatedImagesLoaded().ConfigureAwait(false);

            if (name is null)
            {
                throw new ArgumentNullException(nameof(name));
            }

            DmetrixAssociatedImageInfo? image = FindAssociatedImageInfo(name);
            if (image is null)
            {
                throw new KeyNotFoundException();
            }

            return image.Format;
        }

        private const int BufferSize = 81920;

        public async ValueTask EncodeAssociatedImageAsync(string name, EncodedImageFileFormat format, Stream stream, CancellationToken cancellationToken)
        {
            FileSourceReaderCache? cache = EnsureNotDisposed();
            await EnsureAssociatedImagesLoaded().ConfigureAwait(false);

            if (name is null)
            {
                throw new ArgumentNullException(nameof(name));
            }

            DmetrixAssociatedImageInfo? image = FindAssociatedImageInfo(name);
            if (image is null)
            {
                throw new KeyNotFoundException();
            }

            if (image.Format != format)
            {
                throw new NotSupportedException();
            }

            byte[]? buffer = null;
            FileSourceReader? reader = null;
            try
            {
                reader = await cache.RentAsync(cancellationToken).ConfigureAwait(false);

                long offset = image.Offset;
                int length = image.Length;
                while (length > 0)
                {
                    int readCount = Math.Min(length, BufferSize);
                    ExpandRentedArray(ref buffer, readCount);
                    readCount = Math.Min(buffer.Length, length);
                    int readBytes = await reader.ReadAsync(offset, buffer.AsMemory(0, readCount), cancellationToken).ConfigureAwait(false);
                    if (readBytes != readCount)
                    {
                        throw new InvalidDataException();
                    }
                    await stream.WriteAsync(buffer, 0, readCount, cancellationToken).ConfigureAwait(false);
                    offset += readCount;
                    length -= readCount;
                }
            }
            finally
            {
                if (!(buffer is null))
                {
                    ArrayPool<byte>.Shared.Return(buffer);
                }
                if (!(reader is null))
                {
                    await reader.DisposeAsync().ConfigureAwait(false);
                }
            }
        }

        public async ValueTask EncodeAssociatedImageAsync(string name, EncodedImageFileFormat format, IBufferWriter<byte> writer, CancellationToken cancellationToken)
        {
            FileSourceReaderCache? cache = EnsureNotDisposed();
            await EnsureAssociatedImagesLoaded().ConfigureAwait(false);

            if (name is null)
            {
                throw new ArgumentNullException(nameof(name));
            }

            DmetrixAssociatedImageInfo? image = FindAssociatedImageInfo(name);
            if (image is null)
            {
                throw new KeyNotFoundException();
            }

            if (image.Format != format)
            {
                throw new NotSupportedException();
            }

            FileSourceReader? reader = null;
            try
            {
                reader = await cache.RentAsync(cancellationToken).ConfigureAwait(false);

                long offset = image.Offset;
                int length = image.Length;
                while (length > 0)
                {
                    int readCount = Math.Min(length, BufferSize);
                    Memory<byte> memory = writer.GetMemory(readCount);
                    readCount = Math.Min(memory.Length, length);
                    int readBytes = await reader.ReadAsync(offset, memory.Slice(0, readCount), cancellationToken).ConfigureAwait(false);
                    if (readBytes != readCount)
                    {
                        throw new InvalidDataException();
                    }
                    writer.Advance(readCount);
                    offset += readCount;
                    length -= readCount;
                }
            }
            finally
            {
                if (!(reader is null))
                {
                    await reader.DisposeAsync().ConfigureAwait(false);
                }
            }
        }

        #endregion

        #region Helper

        internal FileSourceReaderCache GetReaderCache()
        {
            return EnsureNotDisposed();
        }

        public int ConvertToDmetrixLayer(int level)
        {
            level = LevelCount - 1 - level;
            return level + _firstValidLayer;
        }

        public long GetTileOffset(DmetrixFileTileIndex tileIndex)
        {
            return _header.HSize + tileIndex.StartPosition;
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        private FileSourceReaderCache EnsureNotDisposed()
        {
            FileSourceReaderCache? source = _source;
            if (source is null)
            {
                ThrowObjectDisposedException();
            }
            return source;
        }

        [DoesNotReturn]
        private static void ThrowObjectDisposedException()
        {
            throw new ObjectDisposedException(nameof(DmetrixWholeSlideImage));
        }

        private static void ExpandRentedArray([NotNull] ref byte[]? buffer, int size)
        {
            if (buffer is null)
            {
                buffer = ArrayPool<byte>.Shared.Rent(size);
                return;
            }
            if (buffer.Length < size)
            {
                ArrayPool<byte>.Shared.Return(buffer);
                buffer = ArrayPool<byte>.Shared.Rent(size);
            }
        }

        #endregion

        private static readonly DeepZoomProfile[] s_profiles = new DeepZoomProfile[] { new DeepZoomProfile(256, 0, EncodedImageFileFormat.Jpeg) };

        public IEnumerable<DeepZoomProfile> GetDeepZoomProfileHints()
        {
            return s_profiles;
        }

        public ValueTask<DeepZoomTileSource?> OpenDeepZoomAsync(DeepZoomProfile profile, CancellationToken cancellationToken = default)
        {
            if (profile.TileSize == 256 && profile.Overlap == 0 && profile.FileFormat == EncodedImageFileFormat.Jpeg)
            {
                return new ValueTask<DeepZoomTileSource?>(new DmetrixDeepZoomTileSource(this));
            }
            return default;
        }
    }
}
